/*
 * Copyright 2021 TiKV Project Authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package com.webank.tikv.region;

import com.google.common.util.concurrent.ThreadFactoryBuilder;
import com.google.protobuf.ByteString;
import com.webank.tikv.PDClient;
import com.webank.tikv.TiConfiguration;
import com.webank.tikv.exception.GrpcException;
import com.webank.tikv.exception.InternalException;
import com.webank.tikv.exception.InvalidStoreException;
import com.webank.tikv.kvproto.Metapb;
import com.webank.tikv.util.BackOffer;
import com.webank.tikv.util.ChannelFactory;
import com.webank.tikv.util.KeyUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.commons.lang3.tuple.Triple;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;

public class RegionManager {

    private static final Logger LOGGER = LogManager.getLogger(RegionManager.class);

    private final RegionCache cache;
    private final PDClient pdClient;
    private final TiConfiguration conf;
    private final ScheduledExecutorService executor;
    private final StoreHealthChecker storeChecker;
    private final ReplicaSelector replicaSelector;

    public RegionManager(TiConfiguration conf, PDClient pdClient, ChannelFactory channelFactory) {
        this.cache = new RegionCache();
        this.pdClient = pdClient;
        this.conf = conf;
        this.replicaSelector = new ReplicaSelector(conf.getZone());

        ThreadFactory threadFactory = new ThreadFactoryBuilder().setNameFormat("tikv-bg-%d").setDaemon(true).build();
        this.executor = Executors.newSingleThreadScheduledExecutor(threadFactory);

        this.storeChecker = new StoreHealthChecker(channelFactory, pdClient, conf);
        long period = conf.getHealthStoreCheckPeriod();
        this.executor.scheduleAtFixedRate(storeChecker, period, period, TimeUnit.MILLISECONDS);
    }

    public TiRegion getRegionByKey(ByteString key, BackOffer backOffer) {
        TiRegion region = cache.getRegionByKey(key);
        if (region != null) {
            return region;
        }
        String hexKey = KeyUtils.bytesToHex(key.toByteArray());
        LOGGER.warn("key:{} not found in region cache", hexKey);
        try {
            Pair<Metapb.Region, Metapb.Peer> pair = pdClient.getRegionByKey(backOffer, key);
            TiRegion tiRegion = createRegion(pair.getLeft(), pair.getRight(), backOffer);
            region = cache.putRegion(tiRegion);
            LOGGER.warn("refresh cache key:{} to region:{} [{} - {}]", hexKey, tiRegion.getId(),
                KeyUtils.bytesToHex(tiRegion.getStartKey().toByteArray()),
                KeyUtils.bytesToHex(tiRegion.getEndKey().toByteArray()));
            return region;
        } catch (Exception e) {
            LOGGER.warn("fail to find region by key:{} exception:{}", hexKey, e.getMessage());
            return null;
        }
    }

    public Triple<TiRegion, TiStore, Metapb.Peer> getRegionStoreByKey(ByteString key, BackOffer backOffer, boolean replicaRead) {
        TiRegion region = getRegionByKey(key, backOffer);
        if (region == null || !region.isValid()) {
            throw new InternalException("Region invalid: " + region);
        }

        Metapb.Peer peer = replicaRead ? region.getNextReplica() : region.getLeader();
        TiStore store = getStoreById(peer.getStoreId(), backOffer);
        return Triple.of(region, store, peer);
    }

    public TiRegion createRegion(Metapb.Region region, Metapb.Peer leader, BackOffer backOffer) {
        if (leader == null) {
            leader = pdClient.getLeaderByRegionId(backOffer, region.getId());
        }
        List<TiStore> stores = new ArrayList<>();
        List<Metapb.Peer> peers = region.getPeersList();
        for (Metapb.Peer peer : peers) {
            try {
                stores.add(getStoreById(peer.getStoreId(), backOffer));
            } catch (Exception e) {
                LOGGER.warn("store:{} not found, ex: {}", peer.getStoreId(), e.toString());
            }
        }
        Metapb.Region meta = Metapb.Region.newBuilder()
            .mergeFrom(region)
            .clearPeers()
            .addAllPeers(peers)
            .build();
        return new TiRegion(conf, meta, leader, peers, stores, replicaSelector);
    }

    public TiStore getStoreById(long id, BackOffer backOffer) {
        TiStore store = cache.getStoreById(id);
        if (store != null) {
            return store;
        }

        try {
            store = new TiStore(pdClient.getStore(backOffer, id));
        } catch (Exception e) {
            LOGGER.warn("failed to get store:{} from pd, ex:{}", id, e.getMessage());
            throw new GrpcException(e);
        }

        // 从pd获取store id失败或者store状态无效时, 从缓存中移除
        Metapb.Store storeMeta = store.getStore();
        if (storeMeta == null) {
            LOGGER.warn("get invalid store info:{}", id);
            throw new InvalidStoreException(id);
        } else if (storeMeta.getState().equals(Metapb.StoreState.Tombstone)) {
            LOGGER.warn("store:{} is tombstone", id);
            throw new InvalidStoreException(id);
        }
        if (cache.putStore(id, store)) {
            storeChecker.scheduleStoreHealthCheck(store);
        }
        return store;
    }

    public List<TiRegion> getAllRegion() {
        return cache.getAllRegion();
    }

    public List<TiStore> getAllStore() {
        return cache.getAllStore();
    }

    public TiRegion updateLeader(TiRegion region, long storeId) {
        if (region.getLeader().getStoreId() == storeId) {
            return region;
        }
        TiRegion newRegion = region.switchPeer(storeId);
        if (newRegion == null) {
            LOGGER.warn("update leader fail, region:{} cannot switch new store:{}, peer:{}",
                region.getId(), storeId, region.getMeta().getPeersList());
            return null;
        }
        cache.updateRegion(region, newRegion);
        return newRegion;
    }

    public void updateStore(TiStore oldStore, TiStore newStore) {
        if (cache.updateStore(oldStore, newStore)) {
            storeChecker.scheduleStoreHealthCheck(newStore);
        }
    }

    public void putRegion(TiRegion region) {
        cache.putRegion(region);
    }

    public void updateRegion(TiRegion expected, TiRegion region) {
        cache.updateRegion(expected, region);
    }

    public void invalidStore(long storeId) {
        cache.invalidStore(storeId);
    }

    public void invalidRegion(TiRegion region) {
        cache.invalidRegion(region);
    }

    public void close() {
        if (executor != null) {
            executor.shutdownNow();
        }
        cache.clearAll();
    }

}
